#include "dicom/storage/decrypt.h"
#include <surbasic/AESHelper.h>
#include <zlib.h>

#include "dicom/storage/encrypt.h"

namespace dicom {

Status Decryptor::DoDecrypt(const std::string& hex_key, ReaderStream* src,
                            WriterStream* out) {
  Status s;

  const int kBlockSize = 4096;
  unsigned char buf[kBlockSize];
  SDBasic::aes::Aes256Decrypt dec(
      SDBasic::bytearray::ByteArray::fromHexStr(hex_key), kBlockSize);

  int ret, flush;
  unsigned have;
  z_stream strm;
  unsigned char outbuf[kBlockSize];
  strm.zalloc = Z_NULL;
  strm.zfree = Z_NULL;
  strm.opaque = Z_NULL;

  ::inflateInit(&strm);

  while (1) {
    int n = src->Read(buf, sizeof(buf));
    if (n < 0) {
      s = Status(Status::kOther, "read failed");
      goto error;
    }

    if (n == 0) {
      if (!dec.Finalize()) {
        s = Status(Status::kOther, "AES256.Finalize() failed");
        goto error;
      }
      flush = Z_FINISH;
    } else {
      if (!dec.Update(buf, n)) {
        s = Status(Status::kOther, "AES256.Update() failed");
        goto error;
      }
      flush = Z_NO_FLUSH;
    }

    strm.avail_in = dec.len();
    strm.next_in = const_cast<Bytef*>(dec.data());

    do {
      strm.avail_out = kBlockSize;
      strm.next_out = outbuf;
      ret = inflate(&strm, flush);
      assert(ret != Z_STREAM_ERROR);
      have = kBlockSize - strm.avail_out;
      switch (ret) {
        case Z_NEED_DICT:
        case Z_DATA_ERROR:
        case Z_MEM_ERROR:
          s = Status(Status::kOther, "zlib::inflate() failed");
          goto error;
      }

      if (out->Write(outbuf, have) < 0) {
        s = Status(Status::kOther, "failed to write output");
        goto error;
      }
    } while (strm.avail_out == 0);

    if (ret == Z_STREAM_END) break;
  }

  ::inflateEnd(&strm);
  return s;

error:
  ::inflateEnd(&strm);
  return s;
}

Status Decryptor::DoDecryptAndHash(const std::string& hex_key,
                                   ReaderStream* src, WriterStream* out,
                                   HashEntry* hash) {
  Status s;

  dicom::MD5Helper src_md5;
  dicom::SHA1Helper src_sha1;
  dicom::SHA1Helper enc_sha1;
  dicom::SHA512Helper src_sha512;

  size_t src_size = 0;
  size_t enc_size = 0;

  const int kBlockSize = 4096;
  unsigned char buf[kBlockSize];
  SDBasic::aes::Aes256Decrypt dec(
      SDBasic::bytearray::ByteArray::fromHexStr(hex_key), kBlockSize);

  int ret, flush;
  unsigned have;
  z_stream strm;
  unsigned char outbuf[kBlockSize];
  strm.zalloc = Z_NULL;
  strm.zfree = Z_NULL;
  strm.opaque = Z_NULL;

  ::inflateInit(&strm);

  while (1) {
    int n = src->Read(buf, sizeof(buf));
    if (n < 0) {
      s = Status(Status::kOther, "read failed");
      goto error;
    }

    if (n == 0) {
      if (!dec.Finalize()) {
        s = Status(Status::kOther, "AES256.Finalize() failed");
        goto error;
      }
      flush = Z_FINISH;
    } else {
      if (!dec.Update(buf, n)) {
        s = Status(Status::kOther, "AES256.Update() failed");
        goto error;
      }
      enc_sha1.Update(buf, n);
      enc_size += n;
      flush = Z_NO_FLUSH;
    }

    strm.avail_in = dec.len();
    strm.next_in = const_cast<Bytef*>(dec.data());

    do {
      strm.avail_out = kBlockSize;
      strm.next_out = outbuf;
      ret = inflate(&strm, flush);
      assert(ret != Z_STREAM_ERROR);
      have = kBlockSize - strm.avail_out;
      switch (ret) {
        case Z_NEED_DICT:
        case Z_DATA_ERROR:
        case Z_MEM_ERROR:
          s = Status(Status::kOther, "zlib::inflate() failed");
          goto error;
      }

      if (out->Write(outbuf, have) < 0) {
        s = Status(Status::kOther, "failed to write output");
        goto error;
      }
      src_md5.Update(outbuf, have);
      src_sha512.Update(outbuf, have);
      src_sha1.Update(outbuf, have);
      src_size += have;
    } while (strm.avail_out == 0);

    if (ret == Z_STREAM_END) break;
  }

  ::inflateEnd(&strm);

  src_md5.Finalize();
  src_sha1.Finalize();
  src_sha512.Finalize();
  enc_sha1.Finalize();

  hash->src_md5 = src_md5.digest();
  hash->src_size = src_size;
  hash->src_digest = StringPrintf("%s%s%lx", src_md5.digest().c_str(),
                                  src_sha1.digest().c_str(), src_size);
  hash->src_sha512 = src_sha512.digest();
  hash->hexkey = hex_key;
  hash->enc_size = enc_size;
  hash->enc_digest = enc_sha1.digest();
  return s;

error:
  ::inflateEnd(&strm);
  return s;
}

Status DecryptAndUnzipFile(const std::string& hexkey,
                           const std::string& src_file,
                           const std::string& dest_file) {
  Status s;

  FileReader reader;
  s = reader.Open(src_file);
  if (!s.ok()) {
    return s;
  }

  FileWriter writer;
  s = writer.Open(dest_file);
  if (!s.ok()) {
    reader.Close();
    return s;
  }

  Decryptor dec;
  s = dec.DoDecrypt(hexkey, &reader, &writer);
  reader.Close();
  writer.Close();
  return s;
}

Status DecryptAndUnzipFileAndHash(const std::string& hexkey,
                                  const std::string& src_file,
                                  const std::string& dest_file,
                                  HashEntry* hash) {
  Status s;

  FileReader reader;
  s = reader.Open(src_file);
  if (!s.ok()) {
    return s;
  }

  FileWriter writer;
  s = writer.Open(dest_file);
  if (!s.ok()) {
    reader.Close();
    return s;
  }

  Decryptor dec;
  s = dec.DoDecryptAndHash(hexkey, &reader, &writer, hash);
  if (s.ok()) {
    assert(hash->IsComplete());
  }
  reader.Close();
  writer.Close();
  return s;
}
}
